package com.openness.crawler.lock;

import java.util.Comparator;
import java.util.List;
import java.util.NavigableSet;
import java.util.TreeSet;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Stat;

import com.openness.crawler.util.Util;
import com.openness.crawler.zookeeper.CountDownLatchWatcher;
import com.openness.crawler.zookeeper.ProtocolSupport;
import com.openness.crawler.zookeeper.ZooKeeperOperation;

/**
 * @author Rain
 * @version 2013-5-21 上午9:26:07
 */

public class DistributedLock extends ProtocolSupport {

	private ZooKeeper zooKeeper;

	private String dir;

	private String prefix;

	private ThreadLocal<String> paths = new ThreadLocal<String>() {

		protected String initialValue() {
			return null;
		};

	};

	private static final String SEPARATOR = "_";

	private static final List<ACL> ACL = ZooDefs.Ids.OPEN_ACL_UNSAFE;

	private static final CreateMode CREATEMODE = CreateMode.EPHEMERAL_SEQUENTIAL;

	private static final int VERSION = -1;

	private static class PathComparator implements Comparator<String> {

		@Override
		public int compare(String path1, String path2) {
			int number1 = Integer.valueOf(path1.substring(path1
					.indexOf(SEPARATOR) + 1));

			int number2 = Integer.valueOf(path2.substring(path2
					.indexOf(SEPARATOR) + 1));

			return number1 == number2 ? 0 : (number1 > number2 ? 1 : -1);
		}

	}

	public DistributedLock(ZooKeeper zooKeeper, String dir, String prefix) {
		this.zooKeeper = zooKeeper;

		this.dir = dir;

		this.prefix = prefix;
	}

	public void lock() {
		String path = paths.get();
		if (path != null && !path.isEmpty()) {
			return;
		}

		try {
			path = retryOperation(new ZooKeeperOperation<String>() {

				private String path = null;

				@Override
				public String execute() throws Exception {
					if (path != null) {
						return path;
					}

					boolean exist = false;

					List<String> child = zooKeeper.getChildren(dir, false);

					for (String children : child) {
						try {
							byte[] data = zooKeeper.getData(dir + "/"
									+ children, false, null);

							if (data.equals(Util.getData())) {
								path = dir + "/" + children;

								exist = true;

								break;
							}
						} catch (KeeperException.NoNodeException e) {
							continue;
						}
					}

					if (!exist) {
						path = zooKeeper.create(dir + "/" + prefix + SEPARATOR,
								Util.getData(), ACL, CREATEMODE);
					}

					return path;
				}

			});
		} catch (Exception e) {
			throw new RuntimeException("DistributedLock lock create: "
					+ e.toString());
		}

		NavigableSet<String> children = new TreeSet<String>(
				new PathComparator());

		try {
			List<String> zkChildren = retryOperation(new ZooKeeperOperation<List<String>>() {

				private List<String> zkChildren = null;

				@Override
				public List<String> execute() throws Exception {
					if (zkChildren != null) {
						return zkChildren;
					}

					zkChildren = zooKeeper.getChildren(dir, false);

					return zkChildren;
				}

			});

			for (String child : zkChildren) {
				children.add(child);
			}
		} catch (Exception e) {
			throw new RuntimeException("DistributedLock lock getChildren: "
					+ e.toString());
		}

		if (children.size() == 0) {
			throw new RuntimeException("Aftre create , ZooKeeper dir: " + dir
					+ "'s children size is 0");
		}

		if (path.equals(dir + "/" + children.first())) {
			paths.set(path);

			return;
		}

		final CountDownLatchWatcher watcher = new CountDownLatchWatcher();

		final String nextLowest = children.lower(path);

		try {
			Stat stat = retryOperation(new ZooKeeperOperation<Stat>() {

				@Override
				public Stat execute() throws Exception {
					return zooKeeper.exists(dir + "/" + nextLowest, watcher);
				}

			});
			boolean exist = stat == null ? false : true;

			if (exist) {
				watcher.await();
			}

			paths.set(path);
		} catch (Exception e) {
			throw new RuntimeException("DistributedLock lock exists: "
					+ e.toString());
		}
	}

	public void unlock() {
		final String path = paths.get();
		if (path == null || path.isEmpty()) {
			return;
		}

		try {
			retryOperation(new ZooKeeperOperation<Object>() {

				@Override
				public Object execute() throws Exception {
					if (zooKeeper.exists(path, false) != null) {
						zooKeeper.delete(path, VERSION);
					}

					return null;
				}

			});

			paths.remove();
		} catch (Exception e) {
			throw new RuntimeException("DistributedLock unlock delete: "
					+ e.toString());
		}
	}

}
